import glob
import os.path
import pickle

import matplotlib.pyplot as plt
import numpy as np
import scipy.io.wavfile as wav
from cut_wave import *
import cv2
import copy

from refactor_code.utils import *



class SingleWavShower:
    def __init__(self, filePath, params):
        self.filePath = filePath
        self.params = params
        self.sonicFreq = params.sonicFreq
        self.lipOffset = params.lipOffset
        self.nfft = params.nfft
        self.frameTime = params.frameTime
        self.aheadTime = params.aheadTime
        self.isGetRange = params.isGetRange
        self.savePicDir = params.savePicDir

    def showOneWav(self):
        fs, originSignal = wav.read(self.filePath)
        sigLen = len(originSignal)
        time = sigLen/fs

        xData = np.linspace(0, time, sigLen, endpoint=False)
        yData = originSignal

        xLabel = 'Time(s)'
        yLabel = 'Amplitude'

        plt.plot(xData, yData)
        plt.title(self.filePath)
        plt.xlabel(xLabel)
        plt.ylabel(yLabel)

        innerPath = getInnerDir(self.filePath)
        saveDir = os.path.join(self.savePicDir, innerPath, 'wav')
        savePltPic(saveDir, self.filePath)

        plt.show()



    def showOneFreq(self, sonicFreq):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        if self.isGetRange:
            borderFreq = self.getBorderFreq(sonicFreq, self.lipOffset)
            data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        plt.xlabel('Time(s)')
        plt.ylabel('Freqency(Hz)')
        plt.pcolormesh(t, f, Zxx)
        plt.show()

    def showFreqWithRange(self, start, end):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        data = self.getRangeSTFTData(data, [start, end])
        f, t, Zxx = data

        plt.title(self.filePath)
        plt.xlabel('Time(s)')
        plt.ylabel('Freqency(Hz)')
        plt.pcolormesh(t, f, Zxx)
        plt.show()


    def showSignalGradient(self, sonicFreq):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        plt.title(self.filePath)
        plt.xlabel('Time(s)')
        plt.ylabel('Freqency(Hz)')
        plt.pcolormesh(t, f, Zxx)
        plt.show()

    def showMultiFreq(self, startFreq, freqInter, freqNum):
        pass

    def getRangeSTFTData(self, data, borderFreq):
        start = borderFreq[0]
        end = borderFreq[1]
        limits = getLipLimitIndex(start, end, self.fs, self.nfft)
        f, t, Zxx = data
        f = f[limits[0]:limits[1]]
        Zxx = Zxx[limits[0]:limits[1], :]
        data[0] = f
        data[2] = Zxx
        return data

    def getBorderFreq(self, sonicFreq, lipOffset):
        start = sonicFreq - lipOffset
        end = sonicFreq + lipOffset
        return [start, end]

    def getSTFTData(self, fs, originSignal):
        nperseg = int(fs * self.frameTime)
        noverlap = int(fs * (self.frameTime - self.aheadTime))
        f, t, Zxx = signal.stft(originSignal, fs, nperseg=nperseg, noverlap=noverlap, nfft=self.nfft)
        Zxx = np.abs(Zxx)
        return [f, t, Zxx]

    def showSumGradient(self, sonicFreq):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr)-padNum]

        guassArr = GaussSmoothZeroPad(sumArr, 20, 2)
        innerPath = getInnerDir(self.filePath)

        saveDir = os.path.join(self.savePicDir, innerPath, 'gauss')
        plt.plot(guassArr)
        savePltPic(saveDir, self.filePath)
        plt.show()

        saveDir = os.path.join(self.savePicDir, innerPath, 'avg')
        avgArr = AvgZeroPad(sumArr, 20)
        plt.plot(avgArr)
        savePltPic(saveDir, self.filePath)
        plt.show()

    def showCropClip(self):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(self.sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr) - padNum]

        avgArr = AvgZeroPad(sumArr, 20)
        avgArr = avgArr / np.max(avgArr)
        thresh = 0.1

        activation = getActivation(avgArr, thresh)
        clipIndex = getClipIndex(activation)
        concatIndex = getConcatIndex(clipIndex, 20)
        newAct = np.zeros(activation.shape[0])
        resIndex = getMaxLenIndex(concatIndex)
        newAct[resIndex[0]:resIndex[1]] = 1

        innerPath = getInnerDir(self.filePath)
        saveDir = os.path.join(self.savePicDir, innerPath, 'final_crop')
        plt.plot(avgArr)
        plt.plot(newAct)
        savePltPic(saveDir, self.filePath)
        plt.show()

    def getFinalCropLen(self):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(self.sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr) - padNum]

        avgArr = AvgZeroPad(sumArr, 20)
        avgArr = avgArr / np.max(avgArr)
        thresh = 0.1

        activation = getActivation(avgArr, thresh)
        clipIndex = getClipIndex(activation)
        concatIndex = getConcatIndex(clipIndex, 20)
        resIndex = getMaxLenIndex(concatIndex)

        return (resIndex[1] - resIndex[0])

    def getFinalCrop(self, lengthDict):
        self.fs, originSignal = wav.read(self.filePath)
        data = self.getSTFTData(self.fs, originSignal)
        borderFreq = self.getBorderFreq(self.sonicFreq, self.lipOffset)
        data = self.getRangeSTFTData(data, borderFreq)
        f, t, Zxx = data
        Zxx = subAdjacent(Zxx)
        t = t[:Zxx.shape[1]]

        padNum = 20
        sumArr = np.sum(Zxx, axis=0)
        sumArr = sumArr ** 2
        sumArr = sumArr[padNum:len(sumArr) - padNum]

        avgArr = AvgZeroPad(sumArr, 20)
        avgArr = avgArr / np.max(avgArr)
        thresh = 0.1

        activation = getActivation(avgArr, thresh)
        clipIndex = getClipIndex(activation)
        concatIndex = getConcatIndex(clipIndex, 20)
        resIndex = getMaxLenIndex(concatIndex)

        meanStd = getMeanAndStdOfWordLen(lengthDict)
        word = getWordFromPath(self.filePath)
        cropLen = getCropLenOfOneWord(meanStd, word)
        finalIndex = getFinalIndexRange(resIndex, cropLen)
        finalIndex[1] = min(finalIndex[1], len(avgArr)-1-padNum)
        return [finalIndex[0]+padNum, finalIndex[1]+padNum]


class MultiWavShower:
    def __init__(self, dirPath, params):
        self.dirPath = dirPath
        self.params = params
        self.allPaths = self.getAllPath(dirPath)

    def getAllPath(self, dirPath):
        allPaths = []
        pattern = os.path.join(dirPath, '*.wav')
        for path in glob.glob(pattern):
            allPaths.append(path)
        return allPaths

    def showAllWav(self):
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            obj.showOneWav()

    def showAllSumGradient(self):
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            obj.showSumGradient(self.params.sonicFreq)

    def showAllCropClip(self):
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            obj.showCropClip()

    def getMaxLenOfEachWord(self):
        lengthDict = {}
        for path in self.allPaths:
            word = getWordFromPath(path)
            lengthDict[word] = []
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            curLen = obj.getFinalCropLen()
            word = getWordFromPath(path)
            lengthDict[word].append(curLen)
        print(lengthDict)
        with open('crop_length.dp', 'wb') as f:
            pickle.dump(lengthDict, f)

    def tempVal(self):
        cropLenPath = 'crop_length.dp'
        if os.path.exists(cropLenPath):
            with open(cropLenPath, 'rb') as f:
                lengthDict = pickle.load(f)
        else:
            raise EOFError
        meanStd = getMeanAndStdOfWordLen(lengthDict)
        for key in meanStd.keys():
            print(key, getCropLenOfOneWord(meanStd, key))

    def saveTimeIndexInFile(self):
        cropLenPath = 'crop_length.dp'
        if os.path.exists(cropLenPath):
            with open(cropLenPath, 'rb') as f:
                lengthDict = pickle.load(f)
        else:
            raise EOFError

        timeMap = {}
        for path in self.allPaths:
            obj = SingleWavShower(path, self.params)
            cropIndex = obj.getFinalCrop(lengthDict)
            timeMap[path] = cropIndex
        saveTxt = 'time_mark.txt'
        saveDp = 'time_mark.dp'
        with open(saveTxt, 'w') as f:
            for key in timeMap.keys():
                msg = f'{key} {timeMap[key][0]} {timeMap[key][1]}\n'
                f.write(msg)


